package com.xiekong.rpc.registry.zookeeper;

import com.xiekong.rpc.constant.Constants;
import com.xiekong.rpc.registry.ServiceRegistry;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Random;

/**
 * 服务注册组件，将服务注册到zookeeper
 * @author XieKong
 */
public class ZookeeperServiceRegistry implements ServiceRegistry {
    private static final Logger logger = LoggerFactory.getLogger(ZookeeperServiceRegistry.class);
    private static CuratorFramework client = null;
    private String registryAddress;

    public ZookeeperServiceRegistry(String registryAddress) {
        this.registryAddress = registryAddress;
        this.init();
    }

    /**
     * 创建zookeeper客户端，并创建服务节点
     */
    private void init() {
        try {
            /**
             * 重试机制，每1000毫秒重试一次，总共3次
             */
            if (client == null) {
                RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 3);
                client = CuratorFrameworkFactory.newClient(this.registryAddress, retryPolicy);
                client.start();
            }

            Stat stat = client.checkExists().forPath(Constants.SERVER_PATH);
            if (stat == null) {
                client.create().forPath(Constants.SERVER_PATH);
            }
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("", e);
        }
    }

    /**
     * 注册服务到zookeeper
     * @param serviceName    服务名称
     * @param serviceAddress 服务地址
     */
    @Override
    public void register(String serviceName, String serviceAddress) {
        /**
         * 注意：如果主节点是临时节点就不能构建子节点。
         */
        try {
            String servicePath = Constants.REGISTER_PATH + serviceName;
            Stat stat = client.checkExists().forPath(servicePath);
            if (stat == null) {
                client.create().forPath(servicePath);
            }
            String serviceAddressPath = servicePath + "/" + serviceAddress;
            stat = client.checkExists().forPath(serviceAddressPath);
            if (stat == null) {
                client.create().withMode(CreateMode.EPHEMERAL).forPath(serviceAddressPath);
            }
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("", e);
        }
    }

    /**
     * 根据服务名称查找服务地址
     * @param serviceName 服务名称
     * @return
     */
    @Override
    public String discover(String serviceName) {
        String address = null;
        try {
            // 获取 service 节点
            String servicePath = Constants.REGISTER_PATH + serviceName;
            Stat stat = client.checkExists().forPath(servicePath);
            if (stat == null) {
                throw new RuntimeException("找不到服务节点 [" + servicePath + "]");
            }
            List<String> addressList = client.getChildren().forPath(servicePath);
            if (addressList.size() == 1) {
                address = addressList.get(0);
            } else {
                // 若存在多个地址，则随机获取一个地址
                Random random = new Random();
                address = addressList.get(random.nextInt(addressList.size()));
            }
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("", e);
        }
        return address;
    }
}
